package com.ruoyi.framework.websocket;


import com.volcengine.ark.runtime.model.completion.chat.*;
import com.volcengine.ark.runtime.service.ArkService;
import io.reactivex.Flowable;
import io.reactivex.disposables.Disposable;
import okhttp3.ConnectionPool;
import okhttp3.Dispatcher;

import javax.websocket.Session;
import java.util.ArrayList;
import java.util.List;
import java.util.concurrent.TimeUnit;



public class ArticleAIServer {
    // 从环境变量中获取您的 API Key。此为默认方式，您可根据需要进行修改
    static String apiKey = "f3447899-2b7c-45cc-afe1-bb9e90ed9b15";
    // 此为默认路径，您可根据业务所在地域进行配置
    static String baseUrl = "https://ark.cn-beijing.volces.com/api/v3";
    static ConnectionPool connectionPool = new ConnectionPool(5, 1, TimeUnit.SECONDS);
    static Dispatcher dispatcher = new Dispatcher();
    static ArkService service = ArkService.builder().dispatcher(dispatcher).connectionPool(connectionPool).baseUrl(baseUrl).apiKey(apiKey).build();


    public void  aiDialogueStream( String userMsg, Session session) {
        List<ChatMessage> messagesList = new ArrayList<>();

        messagesList.add(
                ChatMessage.builder().role(ChatMessageRole.USER)
                        .content(userMsg)
                        .build()
        );

        ChatCompletionRequest chatCompletionRequest = ChatCompletionRequest.builder()
                .model("ep-20250308142727-jkxrv")
                .messages(messagesList)
                .build();


        Flowable<ChatCompletionChunk> stream = service.streamChatCompletion(chatCompletionRequest);

        // 订阅流并处理数据
        Disposable disposable = stream
                .doOnError(Throwable::printStackTrace)
                .subscribe(
                        partialResult -> {
                            String content = (String) partialResult.getChoices().get(0).getMessage().getContent();
                            WebSocketUsers.sendMessageToUserByText(session, content);
                        },
                        error -> {
                            System.err.println("发生错误: " + error.getMessage());
                        },
                        () -> {
                            System.out.println("\n流式响应完成");
                        }
                );

        // 等待流处理完成
        try {
            while (!disposable.isDisposed()) {
                Thread.sleep(100);
            }
        } catch (InterruptedException e) {
            e.printStackTrace();
        }

        // 关闭服务的执行器
        //service.shutdownExecutor();
    }
}
